from trainer.abs_trainer import AbsTrainer
import os.path as op
from eval import SimModel, LSD
import torch


class Trainer(AbsTrainer):
    def __init__(
        self,
        model,
        tr_data,
        cv_data,
        optim,
        config,
        ckpt_path,
        device,
        rank,
        logger,
    ):
        super().__init__(
            model,
            tr_data,
            cv_data,
            optim,
            config,
            ckpt_path,
            device,
            rank,
            logger,
        )
        print(f"using trainer at {op.abspath(__file__)}")
        self.simModel = SimModel(self.device, 16000, "english")
        self.lsd = LSD()
        self.lsd.to(device)

    def get_res(self, res, emb, total_loss, mse_loss, ssim_loss):
        result = {}
        audio_hat = self.model.module.recon(res)
        audio = self.model.module.recon(emb)
        lsd_loss = self.lsd.compute_lsd_batch(audio_hat, audio)
        sim = self.simModel.sim(audio, audio_hat)
        result["lsd"] = torch.tensor(lsd_loss, device=audio.device)
        result["sim"] = torch.tensor(sim, device=audio.device)
        result["loss"] = total_loss
        result["mse_loss"] = mse_loss
        result["ssim_loss"] = ssim_loss
        return result

    def _train_one_batch(self, batch, data, optim, if_log) -> dict:
        data = data[0].to(self.device)  # [B,T,E]
        res, emb, total_loss, mse_loss, ssim_loss = self.model(data)
        total_loss.backward()

        optim.step()
        optim.zero_grad()
        if if_log:
            return self.get_res(res, emb, total_loss, mse_loss, ssim_loss)
        return None

    def _eval_one_batch(self, data) -> dict:
        data = data[0].to(self.device)  # [B,T,E]
        res, emb, total_loss, mse_loss, ssim_loss = self.model(data)
        return self.get_res(res, emb, total_loss, mse_loss, ssim_loss)
